#!/usr/bin/env python

import numpy as N
import glob

import Biskit.tools as T
from Biskit import PDBModel, PDBDope

def addSideChainExposure( m ):
    """
    Convert per-atom accessibility into side chain accessibility per residue
    """
    mask_sc = N.logical_not( m.maskBB() ) ## all side chain atoms

    ## mark residues with any side chain atom (all but GLY)
    resmask = m.atom2resMask( mask_sc )

    ## average relAS over all side chain atoms of each residue
    m_sc = m.compress( mask_sc )  ## model with only side chain atoms, lacks GLY
    scASA= m_sc.atom2resProfile( 'relAS', N.average ) ## now per residue

    ## put side chain average into residue profile, 0 for Glycines
    m.residues.set( 'scAS', scASA, mask=resmask, default=0 )

def exposedO( m ):
    """
    Mark exposed terminal oxygens of SER, THR, TYR 
    """
    m_o = m.mask( ['OG', 'OG1', 'OH'] )
    return N.logical_and( m['relAS'] > 10., m_o )
    

def reportPdb( m, fname ):
    """
    create PDB file with side chain exposure mapped to B-Factor column and
    exposed S/T/Y Oxygen labelled in occupancy column.
    """
    m['temperature_factor'] = m.res2atomProfile('scAS')
    m['occupancy'] = N.logical_not(m['exo']) * 1. ## 0 means exposed O

    m.writePdb( fname )


def tabify( l ):
    """convert list or array to tab-delimited string."""
    return '\t'.join( [str(x) for x in l] )


def reportExposed( m, fname ):
    """Report exposed S/T/Y in a text file"""
    try:
        f = open(fname, 'w')
        f.write(
            '# (1) PDB sequence, (2) exposed S/T/Y (position 0...N), (3) all S/T/Y\n')
        
        f.write( m.sequence() + '\n' )

        mask_exo = m.atom2resProfile( m['exo'], N.sum )
        f.write( tabify( N.flatnonzero( mask_exo ) )  + '\n')

        mask_sty = N.array( [ s in ['S','T','Y'] for s in m.sequence() ] )
        f.write( tabify( N.flatnonzero( mask_sty ) ) + '\n')
        
    except IOError, why:
        print 'error creating report file %s:\n %r' % (fname, why) 

####### MAIN ############

## fname = 'Q99383.6.pdb'

fnames = glob.glob( '*.pdb' )

for f in fnames:

    print 'processing %s...' % f,
    try:
        m = PDBModel( f )

        m['modscore'] = m['temperature_factor']  ## copy quality score, perhaps useful
        m['aa'] = m.sequence() ## create residue profile with single-letter AA 

        d = PDBDope( m )
        d.addSurfaceRacer()  ## accessibilities are now in atom profile m['relAS']

        addSideChainExposure( m )  ## side chain exposure now in residue profile 'scAS'
        m['exo'] = exposedO( m )   ## exposed T/Y/S oxygen now marked in atom profile 'exo'

        m.saveAs( T.stripFilename( f ) + '.model' ) ## save object to disc

        reportExposed( m, T.stripFilename( f ) + '_exposure.tab' )

        print 'OK'

    except:
        print 'Error: ' + T.lastError()

##     reportPdb(     m, T.stripFilename( fname ) + '_exposure.pdb' )
